#!/bin/bash

# ==============================================================================
# SCRIPT TO CONFIGURE MULTIPLE CONTAINERD REGISTRY MIRRORS
#
# This script automates the setup of registry mirrors for multiple registries
# using the modern "hosts.d" directory pattern. It is designed to be
# idempotent and easily extensible.
#
# It will configure:
# - docker.io mirrors
# - registry.k8s.io mirrors
#
# PREREQUISITES:
# - Run this script as a user with sudo privileges.
# - The target system must be using containerd as its container runtime.
# ==============================================================================

# Exit immediately if a command exits with a non-zero status.
set -e

# ==============================================================================
# VARIABLES
#
# Define the registry-to-mirror mappings in an associative array.
# Add new registries and their mirrors here.
# ==============================================================================

# Define the containerd configuration path where host files will be stored.
CONFIG_PATH="/etc/containerd/certs.d"

declare -A MIRROR_MAP
# Mirrors for Docker Hub.
MIRROR_MAP["docker.io"]="https://docker.m.daocloud.io https://docker.mirrors.ustc.edu.cn"

# Mirrors for Google's k8s registry.
# The `registry.k8s.io` is a redirector, so mirrors are less common than for Docker.
# This example uses a common public mirror.
MIRROR_MAP["registry.k8s.io"]="https://registry.aliyuncs.com/google_containers"

# ==============================================================================
# FUNCTIONS
# ==============================================================================

# Function to check for required commands.
check_dependencies() {
    if ! command -v tee &> /dev/null; then
        echo "Error: 'tee' command not found. Please install it."
        exit 1
    fi
}

# Function to add the config_path to containerd's main config.toml.
configure_main_config() {
    echo "Checking and updating /etc/containerd/config.toml..."
    # Check if the config_path is already set. If not, add it.
    if ! grep -q "config_path = \"${CONFIG_PATH}\"" /etc/containerd/config.toml; then
        echo "Adding config_path to /etc/containerd/config.toml."
        sudo sed -i '/^\[plugins\."io.containerd.grpc.v1.cri"\.registry\]/a\  config_path = "'"${CONFIG_PATH}"'"' /etc/containerd/config.toml
    else
        echo "/etc/containerd/config.toml is already configured."
    fi
}

# Function to create the hosts.toml file for a specific registry.
create_hosts_file() {
    local registry=$1
    local mirrors=$2

    echo "Creating directory and hosts.toml for ${registry}..."
    local hosts_dir="${CONFIG_PATH}/${registry}"
    sudo mkdir -p "${hosts_dir}"
    local hosts_file="${hosts_dir}/hosts.toml"

    # Start the hosts.toml content. The default server is the upstream one.
    local hosts_content="server = \"https://registry-1.${registry}\""

    # Handle special case for registry.k8s.io which doesn't have a registry-1 subdomain.
    if [[ "${registry}" == "registry.k8s.io" ]]; then
        hosts_content="server = \"https://registry.k8s.io\""
    fi

    # Add each mirror to the hosts.toml content.
    # The read -a command splits the string of mirrors into an array.
    read -ra mirror_array <<< "$mirrors"
    for mirror in "${mirror_array[@]}"; do
        hosts_content+="\n\n[host.\"${mirror}\"]\n  capabilities = [\"pull\", \"resolve\"]"
    done

    # Write the content to the hosts.toml file.
    # The -e option for echo allows backslash escapes to be interpreted.
    echo -e "${hosts_content}" | sudo tee "${hosts_file}" > /dev/null
    echo "Successfully created ${hosts_file}."
}

# Function to restart containerd to apply the new configuration.
restart_containerd() {
    echo "Restarting containerd service..."
    sudo systemctl daemon-reload
    sudo systemctl restart containerd
    echo "Containerd service restarted. Verification is recommended."
}

# ==============================================================================
# MAIN EXECUTION
# ==============================================================================

echo "Starting containerd registry mirror configuration script."

# 1. Check for required system commands.
check_dependencies

# 2. Add config_path to the main containerd config.
configure_main_config

# 3. Loop through the associative array and create hosts.toml for each registry.
for registry in "${!MIRROR_MAP[@]}"; do
    create_hosts_file "${registry}" "${MIRROR_MAP[${registry}]}"
done

# 4. Restart the containerd service.
restart_containerd

echo "Script completed successfully. All configured registries now use mirrors."

# Exit with a success code.
exit 0
